"""
StarGAN v2
Copyright (c) 2020-present NAVER Corp.

This work is licensed under the Creative Commons Attribution-NonCommercial
4.0 International License. To view a copy of this license, visit
http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""

import os
import pickle

import paddle
from paddle import fluid

from utils.logger import get_logger
logger = get_logger(__name__)


class CheckpointIO(object):
    def __init__(self, fname_template, **kwargs):
        os.makedirs(os.path.dirname(fname_template), exist_ok=True)
        self.fname_template = fname_template
        self.module_dict = kwargs
        # logger.debug(f'CheckpointIO __init__    module_dict: {self.module_dict.keys()}')

    def register(self, **kwargs):
        self.module_dict.update(kwargs)

    def save(self, step):
        fname = self.fname_template.format(step)
        for name, module in self.module_dict.items():
            if fname.endswith('_nets.ckpt'):
                final_fname = fname.replace('_nets.ckpt', f'_{name}')
            elif fname.endswith('_nets_ema.ckpt'):
                final_fname = fname.replace('_nets_ema.ckpt', f'_{name}_ema')
            elif fname.endswith('_optims.ckpt'):
                final_fname = fname.replace('_optims.ckpt', f'_{name}')
            logger.info('Saving checkpoint into %s...' % final_fname)
            fluid.save_dygraph(module.state_dict(), final_fname)

    def load(self, step):
        fname = self.fname_template.format(step)

        logger.info(f'Load: {fname}')
        state_dicts = pickle.load(open(fname, 'rb'))

        for name, module in self.module_dict.items():
            module.set_dict(state_dicts[name])
